import qrcode
from qrcode.exceptions import DataOverflowError
from io import BytesIO

class QRCodeGenerator:
    def __init__(self, version=1, error_correction=qrcode.constants.ERROR_CORRECT_M, box_size=10, border=4):
        """
        初始化二维码生成器
        :param version: 二维码版本 (1-40)
        :param error_correction: 纠错等级
        :param box_size: 每个小格子像素数
        :param border: 二维码边距
        """
        self.version = version
        self.error_correction = error_correction
        self.box_size = box_size
        self.border = border
        self.qr_bytes = None

    def generate(self, data, fill_color="black", back_color="white", output_format="PNG"):
        """
        生成二维码
        :param data: 要编码的数据
        :param fill_color: 二维码颜色
        :param back_color: 背景颜色
        :param output_format: 输出格式 (PNG/SVG)
        :return: 二维码字节数据
        """
        if not data:
            raise ValueError("输入数据不能为空")

        try:
            qr = qrcode.QRCode(
                version=self.version,
                error_correction=self.error_correction,
                box_size=self.box_size,
                border=self.border,
            )
            qr.add_data(data)
            qr.make(fit=True)

            if output_format.upper() == "SVG":
                from qrcode.image.svg import SvgImage
                img = qr.make_image(image_factory=SvgImage)
                return img.to_string()
            else:
                from PIL import Image
                img = qr.make_image(fill_color=fill_color, back_color=back_color)
                buffer = BytesIO()
                img.save(buffer, format="PNG")
                png_data = buffer.getvalue()
                if len(png_data) == 0:
                    raise ValueError("生成的PNG数据为空")
                self.qr_bytes = png_data
                return png_data

        except DataOverflowError:
            raise ValueError("输入数据超出当前版本容量限制，请选择更高版本")

    def download(self, file_path):
        with open(file_path, "wb") as f:
            f.write(self.qr_bytes)
            

if __name__ == "__main__":
    generator = QRCodeGenerator()
    qr_data = "https://www.baidu.com"  # 替换为你的数据
    qr_bytes = generator.generate(qr_data)
    import os
    os.makedirs("./src/temp/photos", exist_ok=True)
    generator.download("./src/temp/photos/qrcode.png")
    print(f"文件已成功写入")
    print("二维码已生成并保存为 qrcode.png")